import os
from litellm import completion
import litellm

# 启用LiteLLM调试模式
litellm.set_verbose = True

# OPENAI_API_KEY
os.environ["OPENAI_API_KEY"] = "sk-giVozfIOM0quMiug0RowM2sx7bTDAKmhnHJ4D026ulyxIkn0"

base_url="https://poloai.club/v1"
custom_llm_provider="openai"

# 定义函数调用
tools = [
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "获取指定城市的天气信息",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {
                        "type": "string",
                        "description": "城市名称"
                    },
                    "date": {
                        "type": "string", 
                        "description": "日期，格式为YYYY-MM-DD"
                    }
                },
                "required": ["city"]
            }
        }
    }
]

# 打印请求参数
print("Request Parameters:")
print(f"Model: claude-3-7-sonnet-20250219")
print(f"API Base: {base_url}")
print(f"Provider: {custom_llm_provider}")
print(f"Tools: {tools}")
print("----------------------------")

# 发送请求
response = completion(
    model="claude-3-5-sonnet-20241022",
    # model="claude-3-7-sonnet-20250219",
    messages=[
        {"role": "user", "content": "使用我提供给你的tools:get_weather 来回答我北京今天天气怎么样?"}
    ],
    tools=tools,
    tool_choice="auto",
    api_base=base_url,
    custom_llm_provider=custom_llm_provider
)

print("Response:", response)

# 检查是否调用了函数
if response.choices[0].message.get("tool_calls"):
    tool_call = response.choices[0].message["tool_calls"][0]
    print("\nFunction called:", tool_call.function.name)
    print("Arguments:", tool_call.function.arguments)
else:
    print("\nNo function was called")
